#!/usr/bin/env python
# coding:utf-8

import argparse
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
import paddle
from paddle3d.models import pillar_encoder
from paddle3d.models import voxelizers
from paddle3d.models.middle_encoders import pillar_scatter
# from vis_utils import hardvoxelize

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--radar_file', default='/data/zw_1128/tracking_03_select/radar_bin',
                        type=str, help='The radar path.')
    parser.add_argument('--lidar_file', default='/data/zw_1128/tracking_03_select/lidar_bin',
                        type=str, help='The radar path.')
    parser.add_argument('--save_dir', default='/data/zw_1128/tracking_03_select/pseudo_image',
                        type=str, help='The save path.')
    parser.add_argument('--voxel_size', default=[0.05, 0.05, 10], type=list, help='voxel size.')
    parser.add_argument('--point_cloud_range', default=[0, -50, -2, 100.0, 50.0, 8], type=list, help='point cloud range.')
    parser.add_argument('--max_num_points_in_voxel', default=4, type=int, help='point cloud range.')
    parser.add_argument('--max_num_voxels', default=40000, type=int, help='max num voxels.')
    return parser.parse_args()


def pts2bev(pts):
    side_range = (-50, 50)
    fwd_range = (0, 100)
    height_range = (-2, 10)
    x, y, z = pts[:, 0], pts[:, 1], pts[:, 2]

    f_filter = np.logical_and(x > fwd_range[0], x < fwd_range[1])
    s_filter = np.logical_and(y > side_range[0], y < side_range[1])
    h_filter = np.logical_and(z > height_range[0], z < height_range[1])
    filter = np.logical_and(f_filter, s_filter)
    filter = np.logical_and(filter, h_filter)
    indices = np.argwhere(filter).flatten()
    x, y, z = x[indices], y[indices], z[indices]

    res = 0.25
    x_img = (-y / res).astype(np.int32)
    y_img = (-x / res).astype(np.int32)
    x_img = x_img - int(np.floor(side_range[0]) / res)
    y_img = y_img + int(np.floor(fwd_range[1]) / res)

    pixel_value = [255, 255, 255]

    x_max = int((side_range[1] - side_range[0]) / res) + 1
    y_max = int((fwd_range[1] - fwd_range[0]) / res) + 1

    im = np.zeros([y_max, x_max, 3], dtype=np.uint8)
    im[y_img, x_img] = pixel_value

    return im[:, :]


if __name__ == '__main__':
    args = parse_args()
    max_num_points_in_voxel = args.max_num_points_in_voxel
    voxel_size = args.voxel_size
    point_cloud_range = args.point_cloud_range
    H = int((point_cloud_range[3]-point_cloud_range[0])/voxel_size[0])
    W = int((point_cloud_range[4]-point_cloud_range[1])/voxel_size[1])

    encoder = pillar_encoder.PillarFeatureNet(
        feat_channels=(3, ),
        max_num_points_in_voxel=args.max_num_points_in_voxel,
        voxel_size=args.voxel_size,
        point_cloud_range=args.point_cloud_range
    )
    voxelizer = voxelizers.HardVoxelizer(
        voxel_size=args.voxel_size,
        point_cloud_range=args.point_cloud_range,
        max_num_points_in_voxel=args.max_num_points_in_voxel,
        max_num_voxels=args.max_num_voxels
    )
    middle_encoder = pillar_scatter.PointPillarsScatter(
        in_channels=3,
        voxel_size=args.voxel_size,
        point_cloud_range=args.point_cloud_range
    )

    from pathlib import Path
    if Path(args.radar_file).is_file():
        radar_file_list = [args.radar_file]
    else:
        radar_file_list = list(Path(args.radar_file).glob('*.bin'))
        radar_file_list.sort()
    if Path(args.lidar_file).is_file():
        lidar_file_list = [args.lidar_file]
    else:
        lidar_file_list = list(Path(args.lidar_file).glob('*.bin'))
        lidar_file_list.sort()

    R = np.array([[0.996195, -0.0871557, 0],
                 [0.0871557, 0.996195, 0],
                 [0, 0, 1]])
    P = np.array([3.8, 0, 0.5])
    R = np.linalg.inv(R.T)

    for radar_file, lidar_file in zip(radar_file_list, lidar_file_list):
        radar_scan = np.fromfile(radar_file, dtype=np.float32)
        lidar_scan = np.fromfile(lidar_file, dtype=np.float32)
        radar_velo = radar_scan.reshape((-1, 8))
        radar_velo = radar_velo[:, :4]
        lidar_velo = lidar_scan.reshape((-1, 4))
        radar_velo[:, :3] = np.dot((radar_velo[:, :3]-P), R)

        # paddle的体素化
        radar_voxel, radar_coors, num_points_radar_voxel = voxelizer(
            paddle.unsqueeze(paddle.to_tensor(radar_velo), axis=0))
        lidar_voxel, lidar_coors, num_points_lidar_voxel = voxelizer(
            paddle.unsqueeze(paddle.to_tensor(lidar_velo), axis=0))

        # vis_util的体素化
        # voxels, coords, num_points_per_voxel = hardvoxelize(
        #     radar_velo, point_cloud_range, voxel_size, max_num_points_in_voxel,
        #     args.max_num_voxels)

        # 直接把体素化后的二维坐标变为图像
        image_radar = np.zeros((H, W, 3), dtype=np.uint8)
        image_lidar = np.zeros((H, W, 3), dtype=np.uint8)
        coors_radar_array = radar_coors.numpy()
        coors_lidar_array = lidar_coors.numpy()
        for i, j in coors_lidar_array[:, 2:4]:
            image_lidar[j, i] = [255, 255, 255]
        for i, j in coors_radar_array[:, 2:4]:
            image_radar[j, i] = [0, 255, 0]

        # 膨胀+混合处理
        kernel = np.ones((5, 5), np.uint8)
        image_dilated = cv2.dilate(image_radar, kernel, iterations=1)
        image_merge = cv2.addWeighted(image_dilated, 0.5, image_lidar, 0.5, 0)
        # image_merge = cv2.flip(image_merge, 0)

        # Pointnet提取特征为伪图像
        # pillar_feature = encoder(radar_voxel, num_points_radar_voxel, radar_coors)
        # pseudo_image = middle_encoder(pillar_feature, radar_coors, 1)
        # image = paddle.squeeze(pseudo_image).numpy()
        # image = image.transpose(2, 1, 0)
        # image = cv2.flip(image, 0)
        # kernel = np.ones((5, 5), np.uint8)
        # image_dilated = cv2.dilate(image, kernel, iterations=1)

        # # 图像显示+保存
        # cv2.namedWindow('pseudo_image', cv2.WINDOW_NORMAL)
        # cv2.imshow('pseudo_image', image_merge)
        # cv2.waitKey(0)
        # if not os.path.exists(args.save_dir):
        #     os.mkdir(args.save_dir)
        # filename, _ = os.path.splitext(os.path.basename(radar_file))
        # cv2.imwrite(os.path.join(args.save_dir, filename+'.png'), image_merge)
        # # bev_im = pts2bev(radar_velo)
        # # plt.figure('show_pillar')
        # # plt.imshow(image)
        # # plt.show()
        # # input()
        # # cv2.imshow('bev', bev_im)
        # cv2.waitKey(0)
    cv2.destroyAllWindows()
